import random
from regulation.reposition.outlier_detection import anormal_region_detection
from env import  Region_info, NUM_REGION
from typing import List, Dict, Any, Set, Tuple
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

prev_outliers = None

def random_region_point(region, roadmap):

    boundary = Region_info[region]

    x = random.randint(min([boundary[0][0], boundary[-1][0]]), max([boundary[0][0], boundary[-1][0]]))
    y = random.randint(min([boundary[0][1], boundary[-1][1]]), max([boundary[0][1], boundary[-1][1]]))
    # point = (x, y)
    while (x, y) not in roadmap[0]:
      x = random.randint(min([boundary[0][0], boundary[-1][0]]), max([boundary[0][0], boundary[-1][0]]))
      y = random.randint(min([boundary[0][1], boundary[-1][1]]), max([boundary[0][1], boundary[-1][1]]))
    return (x, y)

def random_reposition(repo_observ: Dict[str, Any], roadmap: List):
    reposition = []  # type: List[Dict[str, str]]
    # rids = set([r["driver_id" for r in repo_observ]])

    for r in repo_observ["rider_id"]:
        reposition.append(dict(rider_id = r, destination = random.choice(roadmap[0])))
        print(len(roadmap), random.choice(roadmap[0]))
    return reposition

def useOutlier(repo_observ: Dict[str, Any], model_r, model_int, roadmap):
    global prev_outliers
    reposition = []  # type: List[Dict[str, str]]
    supply_exceeds_demand = list()
    supply_donot_meet_demand = list()
    outliers, model_r, model_int = anormal_region_detection(repo_observ, model_r, model_int)
    regions_data = repo_observ["data"]
    region_infos = list()
    #print(regions_data)
    for i in range(NUM_REGION):
        if outliers[-1][i] == 1:
            region_info = dict(
                    id=i,
                    numIdleR=regions_data[-1][i * 4 + 1],
                    numbusyR=regions_data[-1][i * 4],
                    numIdleO=regions_data[-1][i * 4 + 2]
                )
            region_infos.append(region_info)
    num_rep = np.zeros(10)
    for region in region_infos:
        if region['numIdleR'] + region['numbusyR'] > region['numIdleO']:
            supply_donot_meet_demand.append(region['id'])  # 供不应求
        else:
            supply_exceeds_demand.append(region['id']) # 供过于求

    for r in repo_observ["rider_infos"]:
        if r["status"] == 2 and r["work_region"] in supply_donot_meet_demand:  # 只有在空闲的时候才能进行调度
            num_rep[r["work_region"]] += 1
            if len(supply_exceeds_demand) > 0:
                # distance = [repo_observ["distance"][r["work_region"]][idx] for idx in supply_exceeds_demand]
                # region = supply_exceeds_demand[np.argmin(distance)] # 调度最近的区域

                region = random.choice(supply_exceeds_demand)
                reposition.append(dict(rider_id = r["id"], destination = random_region_point(region,roadmap)))
                    # agent.route = self.random_region_point(region)
            else:
                random_region = random.randint(0, NUM_REGION-1)
                while random_region == r["work_region"]:
                    random_region = random.randint(0, NUM_REGION-1)
                reposition.append(dict(rider_id = r["id"], destination = random_region_point(random_region,roadmap)))
            
    print("异常点：", reposition, num_rep, supply_donot_meet_demand , supply_exceeds_demand, regions_data.shape, outliers.shape, outliers[-1])
    if any(x != 0 for x in outliers[-1]):
        result = 1
    else:
        result = 0
    prev_outliers = plot(regions_data, outliers, prev_outliers, with_out=0)
    return reposition, model_r, model_int, result, outliers[-1]

# 不处理重调度的异常检测
def useOutlier_without(repo_observ: Dict[str, Any], model_r, model_int):
    global prev_outliers
    outliers, model_r, model_int = anormal_region_detection(repo_observ, model_r, model_int)
    regions_data = repo_observ["data"]


    # 只返回异常检测结果和模型
    print("异常点：", regions_data.shape, outliers.shape, sum(outliers))
    prev_outliers = plot(regions_data, outliers, prev_outliers, with_out=1)
    return outliers, model_r, model_int

# with_out=0代表加重调度的异常检测；1代表不加的异常检测
def plot(regions_data, outliers, prev_outliers=None, with_out=0):
    if prev_outliers is not None and np.array_equal(outliers, prev_outliers):
        print("No new anomalies detected. Skipping plot generation.")
        return prev_outliers

    busy_riders = regions_data[:, ::4]
    free_riders = regions_data[:, 1::4]
    order_counts = regions_data[:, 2::4]
    order_amounts = regions_data[:, 3::4]
    anomaly_counts = np.sum(outliers, axis=0)

    suffix = "with_repositon" if with_out == 0 else "without_repositon"

    plt.figure(figsize=(12, 8))
    sns.heatmap(outliers.T, cmap="YlGnBu", cbar=True, xticklabels=10, yticklabels=1)
    plt.xlabel("Time Slots")
    plt.ylabel("Regions")
    plt.title(f"Heatmap of Anomalies in Different Regions Over Time ({suffix})")
    plt.savefig(f"anomalies_heatmap_{suffix}.png")
    plt.close()

    plt.figure(figsize=(10, 6))
    plt.bar(range(len(anomaly_counts)), anomaly_counts, color='skyblue')
    plt.xlabel("Regions")
    plt.ylabel("Number of Anomalies")
    plt.title(f"Number of Anomalies Detected in Each Region ({suffix})")
    plt.xticks(range(len(anomaly_counts)), [f"Region {i}" for i in range(len(anomaly_counts))])
    plt.savefig(f"anomalies_bar_plot_{suffix}.png")
    plt.close()

    fig, axes = plt.subplots(4, 1, figsize=(15, 20))

    sns.heatmap(busy_riders.T, ax=axes[0], cmap="Reds", cbar=True, xticklabels=10, yticklabels=1)
    axes[0].set_title(f"Busy Riders ({suffix})")

    sns.heatmap(free_riders.T, ax=axes[1], cmap="Blues", cbar=True, xticklabels=10, yticklabels=1)
    axes[1].set_title(f"Free Riders ({suffix})")

    sns.heatmap(order_counts.T, ax=axes[2], cmap="Greens", cbar=True, xticklabels=10, yticklabels=1)
    axes[2].set_title(f"Order Counts ({suffix})")

    sns.heatmap(order_amounts.T, ax=axes[3], cmap="Purples", cbar=True, xticklabels=10, yticklabels=1)
    axes[3].set_title(f"Order Amounts ({suffix})")

    for ax in axes:
        ax.set_xlabel("Time Slots")
        ax.set_ylabel("Regions")

    plt.tight_layout()
    plt.savefig(f"metrics_heatmaps_{suffix}.png")
    plt.close()

    time_slots = np.arange(regions_data.shape[0])
    num_regions = regions_data.shape[1] // 4
    regions = [f"Region {i}" for i in range(num_regions)]

    fig, axes = plt.subplots(4, 1, figsize=(15, 20))

    for i, data, title in zip(range(4), [busy_riders, free_riders, order_counts, order_amounts], 
                              [f"Busy Riders ({suffix})", f"Free Riders ({suffix})", f"Order Counts ({suffix})", f"Order Amounts ({suffix})"]):
        for j in range(data.shape[1]):
            axes[i].plot(time_slots, data[:, j], label=regions[j])
        axes[i].set_title(title)
        axes[i].set_xlabel("Time Slots")
        axes[i].set_ylabel("Count" if i < 3 else "Amount")
        axes[i].legend()

    plt.tight_layout()
    plt.savefig(f"time_series_plots_{suffix}.png")
    plt.close()

    total_busy_riders = np.sum(busy_riders, axis=1)
    total_free_riders = np.sum(free_riders, axis=1)
    total_order_counts = np.sum(order_counts, axis=1)
    total_order_amounts = np.sum(order_amounts, axis=1)

    fig, ax = plt.subplots(figsize=(15, 8))

    time_slots = np.arange(regions_data.shape[0])
    bar_width = 0.35

    ax.bar(time_slots, total_busy_riders, bar_width, label='Busy Riders', color='r')
    ax.bar(time_slots, total_free_riders, bar_width, bottom=total_busy_riders, label='Free Riders', color='b')
    ax.bar(time_slots, total_order_counts, bar_width, bottom=total_busy_riders + total_free_riders, label='Order Counts', color='g')
    ax.bar(time_slots, total_order_amounts, bar_width, bottom=total_busy_riders + total_free_riders + total_order_counts, label='Order Amounts', color='purple')

    ax.set_xlabel('Time Slots')
    ax.set_ylabel('Counts / Amounts')
    ax.set_title(f'Stacked Bar Plot of Different Metrics in Regions Over Time ({suffix})')
    ax.legend()

    plt.tight_layout()
    plt.savefig(f"stacked_bar_plot_{suffix}.png")
    plt.close()

    return outliers